#include <vector>
#include <iostream>
#include <cstdlib>
#include <ctime>
#include <memory>
#include "dense_layer.h"
#include "relu.h"
#include "sigmoid.h"
#include "loss_function.h"

class FNN {
public:
    std::vector<DenseLayer> hiddenLayers;
    DenseLayer outputLayer;
    std::shared_ptr<MatrixOperations> matrixOperations;
    std::shared_ptr<ILossFunction> lossFunction; // 存储损失函数的引用

    /**
     * 构造函数初始化一个全连接神经网络 (FNN)
     * 
     * @param inputSize 输入层的大小，即输入向量的维度。
     * @param hiddenLayerSize 隐藏层的大小，所有隐藏层都将使用这个大小。
     * @param numHiddenLayers 隐藏层的数量，定义了网络中隐藏层的总数。
     * @param outputSize 输出层的大小，即输出向量的维度。
     * @param lr 学习率，用于在训练过程中更新权重。
     * @param matOp 矩阵运算的共享指针，用于执行所有矩阵相关的操作，比如前向传播和反向传播。
     * @param lossFunction 损失函数
     */
    FNN(size_t inputSize, size_t hiddenLayerSize, size_t numHiddenLayers, size_t outputSize, float lr, std::shared_ptr<ActivationFunction> activationFunction, std::shared_ptr<MatrixOperations> matOp, std::shared_ptr<ILossFunction> lossFunction)
    : outputLayer(hiddenLayerSize, outputSize, lr, activationFunction, matOp), matrixOperations(matOp), lossFunction(lossFunction) {
        size_t layerInputSize = inputSize;
        // 根据指定的隐藏层数量创建隐藏层
        for (size_t i = 0; i < numHiddenLayers; ++i) {
            hiddenLayers.emplace_back(DenseLayer(layerInputSize, hiddenLayerSize, lr, std::make_shared<ReLU>(), matOp));
            layerInputSize = hiddenLayerSize; // 下一层的输入大小等于这一层的输出大小
        }
    }

    void train(const std::vector<std::vector<float>>& inputs, const std::vector<std::vector<float>>& expectedOutputs) {
        // 前向传播，计算网络的输出
        auto layerOutputs = inputs;
        for (auto& layer : hiddenLayers) {
            layerOutputs = layer.forward(layerOutputs);
        }
        auto finalOutputs = outputLayer.forward(layerOutputs);

        // 计算输出层的误差
        std::vector<std::vector<float>> outputErrors = lossFunction->gradient(finalOutputs, expectedOutputs);

        // 反向传播
        auto errorGradient = outputLayer.backward(outputErrors);
        // 逆序遍历隐藏层，因为反向传播是从输出层开始的
        for (auto it = hiddenLayers.rbegin(); it != hiddenLayers.rend(); ++it) {
            errorGradient = it->backward(errorGradient);
        }
    }


    std::vector<std::vector<float>> predict(const std::vector<std::vector<float>>& inputs) {
        auto layerOutputs = inputs;
        // 遍历所有隐藏层
        for (auto& layer : hiddenLayers) {
            layerOutputs = layer.forward(layerOutputs);
        }
        // 最后通过输出层
        auto finalOutputs = outputLayer.forward(layerOutputs);
        return finalOutputs;
    }
};
